import pygame
import random
import numpy as np
import math
from enum import Enum

# 初始化pygame
pygame.init()

# 游戏常量
SCREEN_SIZE = 400
GRID_SIZE = 4
GRID_PADDING = 10
TILE_SIZE = (SCREEN_SIZE - (GRID_SIZE + 1) * GRID_PADDING) // GRID_SIZE
BACKGROUND_COLOR = (187, 173, 160)
EMPTY_TILE_COLOR = (205, 193, 180)

# 动画常量
ANIMATION_DURATION = 150  # 移动动画时长（毫秒）
MERGE_ANIMATION_DURATION = 100  # 合并动画时长（毫秒）
EASE_OUT_POWER = 3  # 缓动效果的强度

# 使用支持中文的字体
try:
    FONT = pygame.font.SysFont("SimHei", 40, bold=True)
    SMALL_FONT = pygame.font.SysFont("SimHei", 20, bold=True)
except:
    FONT = pygame.font.SysFont(None, 40, bold=True)
    SMALL_FONT = pygame.font.SysFont(None, 20, bold=True)

# 方块颜色
TILE_COLORS = {
    0: (205, 193, 180),
    2: (238, 228, 218),
    4: (237, 224, 200),
    8: (242, 177, 121),
    16: (245, 149, 99),
    32: (246, 124, 95),
    64: (246, 94, 59),
    128: (237, 207, 114),
    256: (237, 204, 97),
    512: (237, 200, 80),
    1024: (237, 197, 63),
    2048: (237, 194, 46),
    4096: (237, 194, 46),
    8192: (237, 194, 46)
}


class Direction(Enum):
    UP = 0
    RIGHT = 1
    DOWN = 2
    LEFT = 3


class AnimationState(Enum):
    IDLE = 0
    MOVING = 1
    MERGING = 2


class Tile:
    """代表一个方块的类，包含位置和动画信息"""

    def __init__(self, row, col, value):
        self.row = row
        self.col = col
        self.value = value
        self.prev_row = row
        self.prev_col = col
        self.animation_start_time = 0
        self.merge_animation_start_time = 0
        self.is_new = False
        self.is_merged = False
        self.merge_from_tiles = []  # 存储合并来源的方块
        self.should_remove = False  # 标记是否应该在动画后移除

    def start_move_animation(self, new_row, new_col, current_time):
        """开始移动动画"""
        self.prev_row = self.row
        self.prev_col = self.col
        self.row = new_row
        self.col = new_col
        self.animation_start_time = current_time

    def start_merge_animation(self, current_time):
        """开始合并动画"""
        self.merge_animation_start_time = current_time

    def get_animated_position(self, current_time):
        """获取当前动画位置"""
        if self.animation_start_time == 0:
            return self.row, self.col

        elapsed = current_time - self.animation_start_time
        if elapsed >= ANIMATION_DURATION:
            return self.row, self.col

        # 使用ease-out缓动函数
        progress = elapsed / ANIMATION_DURATION
        eased_progress = 1 - math.pow(1 - progress, EASE_OUT_POWER)

        animated_row = self.prev_row + (self.row - self.prev_row) * eased_progress
        animated_col = self.prev_col + (self.col - self.prev_col) * eased_progress

        return animated_row, animated_col

    def get_merge_scale(self, current_time):
        """获取合并动画的缩放比例"""
        if self.merge_animation_start_time == 0:
            return 1.0

        elapsed = current_time - self.merge_animation_start_time
        if elapsed >= MERGE_ANIMATION_DURATION:
            return 1.0

        # 合并动画：先放大再缩小回原来大小
        progress = elapsed / MERGE_ANIMATION_DURATION
        # 使用正弦函数创建弹跳效果
        scale = 1.0 + 0.3 * math.sin(progress * math.pi)
        return scale


class Game2048:
    def __init__(self):
        self.grid = np.zeros((GRID_SIZE, GRID_SIZE), dtype=int)
        self.tiles = {}  # 存储所有活动的tile对象
        self.moving_tiles = []  # 存储正在移动但还没到达目标的方块
        self.score = 0
        self.game_over = False
        self.animation_state = AnimationState.IDLE
        self.animation_start_time = 0

        # 添加初始方块
        self.add_new_tile(animate=False)
        self.add_new_tile(animate=False)

    def add_new_tile(self, animate=True):
        """在随机空位置添加一个新的方块"""
        empty_cells = [(i, j) for i in range(GRID_SIZE) for j in range(GRID_SIZE) if self.grid[i][j] == 0]
        if empty_cells:
            i, j = random.choice(empty_cells)
            value = 2 if random.random() < 0.9 else 4
            self.grid[i][j] = value

            # 创建新的tile对象
            tile = Tile(i, j, value)
            if animate:
                tile.is_new = True
            self.tiles[(i, j)] = tile

    def ease_out_back(self, t):
        """回弹缓动函数"""
        c1 = 1.70158
        c3 = c1 + 1
        return 1 + c3 * math.pow(t - 1, 3) + c1 * math.pow(t - 1, 2)

    def move_left(self):
        """向左移动的核心逻辑"""
        moved = False
        new_tiles = {}
        self.moving_tiles = []
        current_time = pygame.time.get_ticks()

        for i in range(GRID_SIZE):
            # 收集当前行的所有非零元素及其对应的tile
            row_values = []
            row_tiles = []

            for j in range(GRID_SIZE):
                if self.grid[i][j] != 0:
                    row_values.append(self.grid[i][j])
                    if (i, j) in self.tiles:
                        row_tiles.append(self.tiles[(i, j)])
                    else:
                        # 如果没有对应的tile，创建一个
                        tile = Tile(i, j, self.grid[i][j])
                        row_tiles.append(tile)

            # 处理移动和合并
            target_col = 0
            k = 0

            while k < len(row_values):
                if k < len(row_values) - 1 and row_values[k] == row_values[k + 1]:
                    # 合并两个相同的数字
                    merged_value = row_values[k] * 2
                    self.score += merged_value

                    # 创建合并后的tile
                    merged_tile = Tile(i, target_col, merged_value)
                    merged_tile.is_merged = True

                    # 获取两个要合并的源tile
                    tile1 = row_tiles[k]
                    tile2 = row_tiles[k + 1]

                    # 让两个源tile都移动到目标位置
                    if tile1.col != target_col:
                        tile1.start_move_animation(i, target_col, current_time)
                        moved = True
                    if tile2.col != target_col:
                        tile2.start_move_animation(i, target_col, current_time)
                        moved = True

                    # 标记源tile在动画完成后应该被移除
                    tile1.should_remove = True
                    tile2.should_remove = True

                    # 将源tile添加到移动列表中，用于渲染
                    self.moving_tiles.extend([tile1, tile2])

                    # 合并后的tile添加到新位置
                    new_tiles[(i, target_col)] = merged_tile

                    k += 2
                else:
                    # 不合并，直接移动
                    tile = row_tiles[k]

                    # 如果位置发生变化，开始动画
                    if tile.col != target_col:
                        tile.start_move_animation(i, target_col, current_time)
                        moved = True
                    else:
                        # 位置没变，但可能需要更新坐标
                        tile.row = i
                        tile.col = target_col

                    new_tiles[(i, target_col)] = tile
                    k += 1

                target_col += 1

            # 更新网格这一行
            new_row = [0] * GRID_SIZE
            for (row, col), tile in new_tiles.items():
                if row == i:  # 如果是当前行的tile
                    new_row[col] = tile.value

            # 检查这一行是否有移动
            if self.grid[i].tolist() != new_row:
                moved = True

            self.grid[i] = new_row

        # 更新tiles字典
        self.tiles = new_tiles
        return moved

    def move(self, direction):
        """移动方块并合并"""
        if self.animation_state != AnimationState.IDLE:
            return

        # 保存原始状态用于比较
        original_grid = self.grid.copy()
        moved = False

        # 根据方向旋转网格和tiles
        if direction == Direction.LEFT:
            moved = self.move_left()
        elif direction == Direction.RIGHT:
            # 水平翻转，向左移动，再翻转回来
            self.grid = np.flip(self.grid, axis=1)
            self.flip_tiles_horizontal()
            moved = self.move_left()
            self.grid = np.flip(self.grid, axis=1)
            self.flip_tiles_horizontal()
        elif direction == Direction.UP:
            # 转置，向左移动，再转置回来
            self.grid = self.grid.T
            self.transpose_tiles()
            moved = self.move_left()
            self.grid = self.grid.T
            self.transpose_tiles()
        elif direction == Direction.DOWN:
            # 转置，水平翻转，向左移动，再恢复
            self.grid = self.grid.T
            self.transpose_tiles()
            self.grid = np.flip(self.grid, axis=1)
            self.flip_tiles_horizontal()
            moved = self.move_left()
            self.grid = np.flip(self.grid, axis=1)
            self.flip_tiles_horizontal()
            self.grid = self.grid.T
            self.transpose_tiles()

        # 如果有移动，开始动画
        if moved:
            self.animation_state = AnimationState.MOVING
            self.animation_start_time = pygame.time.get_ticks()
            # 设置定时器：先等移动动画完成，再开始合并动画
            pygame.time.set_timer(pygame.USEREVENT + 1, ANIMATION_DURATION)

    def flip_tiles_horizontal(self):
        """水平翻转所有tiles的位置"""
        new_tiles = {}
        for (row, col), tile in self.tiles.items():
            new_col = GRID_SIZE - 1 - col
            tile.row = row
            tile.col = new_col
            # 如果正在动画中，也要调整动画的起始和目标位置
            if tile.animation_start_time > 0:
                tile.prev_col = GRID_SIZE - 1 - tile.prev_col
            new_tiles[(row, new_col)] = tile
        self.tiles = new_tiles

        # 也要翻转移动中的tiles
        for tile in self.moving_tiles:
            tile.col = GRID_SIZE - 1 - tile.col
            if tile.animation_start_time > 0:
                tile.prev_col = GRID_SIZE - 1 - tile.prev_col

    def transpose_tiles(self):
        """转置所有tiles的位置"""
        new_tiles = {}
        for (row, col), tile in self.tiles.items():
            new_row, new_col = col, row
            tile.row = new_row
            tile.col = new_col
            # 如果正在动画中，也要调整动画的起始和目标位置
            if tile.animation_start_time > 0:
                tile.prev_row, tile.prev_col = tile.prev_col, tile.prev_row
            new_tiles[(new_row, new_col)] = tile
        self.tiles = new_tiles

        # 也要转置移动中的tiles
        for tile in self.moving_tiles:
            tile.row, tile.col = tile.col, tile.row
            if tile.animation_start_time > 0:
                tile.prev_row, tile.prev_col = tile.prev_col, tile.prev_row

    def update_animation(self, current_time):
        """更新动画状态"""
        if self.animation_state == AnimationState.IDLE:
            return

        elif self.animation_state == AnimationState.MOVING:
            # 检查移动动画是否完成
            if current_time - self.animation_start_time >= ANIMATION_DURATION:
                # 移动动画完成，开始合并动画
                self.animation_state = AnimationState.MERGING

                # 为所有合并的tile开始合并动画
                for tile in self.tiles.values():
                    if tile.is_merged:
                        tile.start_merge_animation(current_time)

                # 清理移动中的tiles
                self.moving_tiles = []

                # 设置合并动画定时器
                pygame.time.set_timer(pygame.USEREVENT + 2, MERGE_ANIMATION_DURATION)

        elif self.animation_state == AnimationState.MERGING:
            # 检查合并动画是否完成
            if current_time - self.animation_start_time >= ANIMATION_DURATION + MERGE_ANIMATION_DURATION:
                self.animation_state = AnimationState.IDLE

                # 清理动画状态
                for tile in self.tiles.values():
                    tile.animation_start_time = 0
                    tile.merge_animation_start_time = 0
                    tile.is_new = False
                    tile.is_merged = False

    def check_game_over(self):
        """检查游戏是否结束"""
        # 检查是否有空格
        if 0 in self.grid:
            self.game_over = False
            return

        # 检查是否有可合并的相邻方块
        for i in range(GRID_SIZE):
            for j in range(GRID_SIZE):
                if j < GRID_SIZE - 1 and self.grid[i][j] == self.grid[i][j + 1]:
                    self.game_over = False
                    return
                if i < GRID_SIZE - 1 and self.grid[i][j] == self.grid[i + 1][j]:
                    self.game_over = False
                    return

        self.game_over = True

    def draw(self, screen):
        """绘制游戏界面"""
        current_time = pygame.time.get_ticks()
        self.update_animation(current_time)

        # 绘制背景
        screen.fill(BACKGROUND_COLOR)

        # 绘制分数
        score_text = FONT.render(f"分数: {self.score}", True, (119, 110, 101))
        screen.blit(score_text, (10, SCREEN_SIZE + 10))

        # 绘制空的网格背景
        for i in range(GRID_SIZE):
            for j in range(GRID_SIZE):
                x = GRID_PADDING + j * (TILE_SIZE + GRID_PADDING)
                y = GRID_PADDING + i * (TILE_SIZE + GRID_PADDING)
                pygame.draw.rect(screen, EMPTY_TILE_COLOR,
                                 (x, y, TILE_SIZE, TILE_SIZE),
                                 border_radius=5)

        # 先绘制正在移动的tiles（即将被移除的）
        for tile in self.moving_tiles:
            if tile.value == 0:
                continue
            self.draw_tile(screen, tile, current_time)

        # 绘制主要的tiles
        for tile in self.tiles.values():
            if tile.value == 0:
                continue
            self.draw_tile(screen, tile, current_time)

        # 如果游戏结束，显示结束信息
        if self.game_over:
            overlay = pygame.Surface((SCREEN_SIZE, SCREEN_SIZE), pygame.SRCALPHA)
            overlay.fill((238, 228, 218, 128))
            screen.blit(overlay, (0, 0))

            game_over_text = FONT.render("游戏结束!", True, (119, 110, 101))
            restart_text = SMALL_FONT.render("按R键重新开始", True, (119, 110, 101))

            screen.blit(game_over_text, (SCREEN_SIZE // 2 - game_over_text.get_width() // 2,
                                         SCREEN_SIZE // 2 - game_over_text.get_height() // 2 - 20))
            screen.blit(restart_text, (SCREEN_SIZE // 2 - restart_text.get_width() // 2,
                                       SCREEN_SIZE // 2 - restart_text.get_height() // 2 + 20))

    def draw_tile(self, screen, tile, current_time):
        """绘制单个方块"""
        # 获取动画位置
        anim_row, anim_col = tile.get_animated_position(current_time)

        # 计算屏幕位置
        x = GRID_PADDING + anim_col * (TILE_SIZE + GRID_PADDING)
        y = GRID_PADDING + anim_row * (TILE_SIZE + GRID_PADDING)

        # 处理缩放效果
        scale = 1.0

        # 新方块的出现动画
        if tile.is_new and current_time - self.animation_start_time < ANIMATION_DURATION:
            progress = (current_time - self.animation_start_time) / ANIMATION_DURATION
            scale = self.ease_out_back(progress) if progress <= 1.0 else 1.0
            scale = max(0.1, min(1.2, scale))

        # 合并动画的缩放效果
        if tile.is_merged:
            merge_scale = tile.get_merge_scale(current_time)
            scale *= merge_scale

        # 计算缩放后的尺寸和位置
        scaled_size = int(TILE_SIZE * scale)
        scaled_x = x + (TILE_SIZE - scaled_size) // 2
        scaled_y = y + (TILE_SIZE - scaled_size) // 2

        # 绘制方块背景
        color = TILE_COLORS.get(tile.value, (60, 58, 50))
        pygame.draw.rect(screen, color,
                         (scaled_x, scaled_y, scaled_size, scaled_size),
                         border_radius=5)

        # 绘制数字
        text_color = (119, 110, 101) if tile.value < 128 else (249, 246, 242)
        font_size = int(40 * scale)
        if font_size > 0:
            try:
                scaled_font = pygame.font.SysFont("SimHei", font_size, bold=True)
            except:
                scaled_font = pygame.font.SysFont(None, font_size, bold=True)

            text = scaled_font.render(str(tile.value), True, text_color)
            text_rect = text.get_rect(center=(x + TILE_SIZE // 2, y + TILE_SIZE // 2))
            screen.blit(text, text_rect)


def main():
    # 创建游戏窗口
    screen = pygame.display.set_mode((SCREEN_SIZE, SCREEN_SIZE + 50))
    pygame.display.set_caption("2048 - 优化版")
    clock = pygame.time.Clock()

    game = Game2048()
    running = True

    while running:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False

            # 移动动画完成，开始合并动画
            elif event.type == pygame.USEREVENT + 1:
                pygame.time.set_timer(pygame.USEREVENT + 1, 0)  # 停止计时器

            # 合并动画完成，添加新方块
            elif event.type == pygame.USEREVENT + 2:
                pygame.time.set_timer(pygame.USEREVENT + 2, 0)  # 停止计时器
                game.add_new_tile()
                game.check_game_over()

            elif not game.game_over and game.animation_state == AnimationState.IDLE:
                if event.type == pygame.KEYDOWN:
                    if event.key == pygame.K_UP or event.key == pygame.K_w:
                        game.move(Direction.UP)
                    elif event.key == pygame.K_RIGHT or event.key == pygame.K_d:
                        game.move(Direction.RIGHT)
                    elif event.key == pygame.K_DOWN or event.key == pygame.K_s:
                        game.move(Direction.DOWN)
                    elif event.key == pygame.K_LEFT or event.key == pygame.K_a:
                        game.move(Direction.LEFT)

            # 重新开始游戏
            if event.type == pygame.KEYDOWN and event.key == pygame.K_r:
                game = Game2048()

        game.draw(screen)
        pygame.display.flip()
        clock.tick(60)  # 60 FPS

    pygame.quit()


if __name__ == "__main__":
    main()